import ptan
import numpy as np

import torch
import torch.nn as nn
import torch.nn.functional as F

import numpy as np
import math


class NoisyLinear(nn.Module):
  def __init__(self, in_features, out_features, std_init=0.5):
    super(NoisyLinear, self).__init__()
    self.module_name = 'noisy_linear'
    self.in_features = in_features
    self.out_features = out_features
    self.std_init = std_init
    self.weight_mu = nn.Parameter(torch.empty(out_features, in_features))
    self.weight_sigma = nn.Parameter(torch.empty(out_features, in_features))
    self.register_buffer('weight_epsilon', torch.empty(out_features, in_features))
    self.bias_mu = nn.Parameter(torch.empty(out_features))
    self.bias_sigma = nn.Parameter(torch.empty(out_features))
    self.register_buffer('bias_epsilon', torch.empty(out_features))
    self.reset_parameters()
    self.reset_noise()

  def reset_parameters(self):
    mu_range = 1 / math.sqrt(self.in_features)
    self.weight_mu.data.uniform_(-mu_range, mu_range)
    self.weight_sigma.data.fill_(self.std_init / math.sqrt(self.in_features))
    self.bias_mu.data.uniform_(-mu_range, mu_range)
    self.bias_sigma.data.fill_(self.std_init / math.sqrt(self.out_features))

  def _scale_noise(self, size):
    '''
    缩放
    缩放到[-1, 1]之间

    这行代码是在进行噪声缩放操作，具体实现了 Factorised Noisy Networks 中的特殊噪声生成方法。让我详细解释：

  ```python
  x.sign().mul_(x.abs().sqrt_())
  ```

  这行代码可以分解为以下步骤：

  1. `x.sign()` - 获取输入张量 x 的符号（+1 或 -1）
  2. `x.abs()` - 计算输入张量的绝对值
  3. `sqrt_()` - 计算绝对值的平方根
  4. `mul_()` - 将符号与平方根相乘（原位操作）

  ### 作用
  - 这个操作实现了论文 [*Noisy Networks for Exploration*](https://arxiv.org/abs/1706.10295) 中提出的噪声生成方法
  - 生成的噪声遵循特定分布，比普通的高斯噪声更适合探索
  - 输出范围在 [-1, 1] 之间，但分布不是均匀的

  ### 数学表达式
  如果输入是随机变量 ε，输出 f(ε) 的计算公式为：
  ```
  f(ε) = sign(ε) * sqrt(|ε|)
  ```

  这种噪声生成方法有助于：
  - 提供更好的探索性能
  - 保持噪声的有界性
  - 产生适当的探索-利用平衡

  在 NoisyLinear 层中，这个噪声用于权重和偏置的扰动，帮助网络进行更有效的探索。
    '''
    x = torch.randn(size)
    return x.sign().mul_(x.abs().sqrt_())

  def reset_noise(self):
    '''
    重制噪声层

    所有的权重的偏置重置到-1到1之间
    '''
    epsilon_in = self._scale_noise(self.in_features)
    epsilon_out = self._scale_noise(self.out_features)
    self.weight_epsilon.copy_(epsilon_out.ger(epsilon_in))
    self.bias_epsilon.copy_(epsilon_out)

  def forward(self, input):
    if self.training:
      return F.linear(input, self.weight_mu + self.weight_sigma * self.weight_epsilon, self.bias_mu + self.bias_sigma * self.bias_epsilon)
    else:
      return F.linear(input, self.weight_mu, self.bias_mu)




class RainbowDQN(nn.Module):
    '''
    虽然代码没有完全的禁用噪音层，但是在训练阶段会学习噪音，使得在推理阶段时噪音对结果的影响很小
    当然也可以增加禁用噪音的参数
    '''
    def __init__(self, input_shape, n_actions, params):
        super(RainbowDQN, self).__init__()
        self.atoms = params['atoms']
        self.v_min = params['value_min']
        self.v_max = params['value_max']
        self.DELTA_Z = (self.v_max - self.v_min) / (self.atoms - 1)
        self.action_space = n_actions

        self.conv = nn.Sequential(
            nn.Conv2d(input_shape[0], 64, kernel_size=8, stride=4),
            nn.ReLU(),
            nn.Conv2d(64, 128, kernel_size=4, stride=2),
            nn.ReLU(),
            nn.Conv2d(128, 128, kernel_size=3, stride=1),
            nn.ReLU()
        )

        self.conv_output_size = self._get_conv_out(input_shape)
        self.fc_h_v = NoisyLinear(self.conv_output_size, params['hidden_size'], std_init=params['noisy_std'])
        self.fc_z_v = NoisyLinear(params['hidden_size'], params['atoms'], std_init=params['noisy_std']) # 状态的价值分布

        self.fc_h_a = NoisyLinear(self.conv_output_size, params['hidden_size'], std_init=params['noisy_std'])
        self.fc_z_a = NoisyLinear(params['hidden_size'], n_actions * params['atoms'], std_init=params['noisy_std']) # 动作的优势分布

        # 对比学习的头，todo作用
        self.W_h = nn.Parameter(torch.rand(self.conv_output_size, params['hidden_size']))
        self.b_h = nn.Parameter(torch.zeros(params['hidden_size']))
        self.layer_norm_h = nn.LayerNorm(params['hidden_size'])

        self.W_c = nn.Parameter(torch.rand(params['hidden_size'], 128))
        self.b_c = nn.Parameter(torch.zeros(128))
        self.layer_norm_c = nn.LayerNorm(128)

        self.W = nn.Parameter(torch.rand(128, 128)) #

        self.register_buffer("supports", torch.arange(self.v_min, self.v_max+self.DELTA_Z, self.DELTA_Z))
        self.softmax = nn.Softmax(dim=1)

    def _get_conv_out(self, shape):
        o = self.conv(torch.zeros(1, *shape))
        return int(np.prod(o.size()))

    def forward(self, x):
        batch_size = x.size()[0]
        fx = x.float() / 255
        conv_out = self.conv(fx).view(batch_size, -1)
        v = self.fc_z_v(F.relu(self.fc_h_v(conv_out)))  # Value stream
        a = self.fc_z_a(F.relu(self.fc_h_a(conv_out)))  # Advantage stream
        # 查看md文档，因为这里说明的是实现的是Dueling DQN
        v, a = v.view(-1, 1, self.atoms), a.view(-1, self.action_space, self.atoms)
        # 这里v + a实现的是Dueling DQN的公式，也就是状态的价值加上动作的优势，组合成最终的Q值
        # 而- a.mean(1, keepdim=True) 这个懂工作是为了去除动作的平均值，避免过于依赖某个动作
        q = v + a - a.mean(1, keepdim=True)  # Combine streams

        h = torch.matmul(conv_out, self.W_h) + self.b_h # Contrastive head
        h = self.layer_norm_h(h)
        h = F.relu(h)
        h = torch.matmul(h, self.W_c) + self.b_c # Contrastive head
        h = self.layer_norm_c(h)
        
        return q, h
    
    def reset_noise(self):
        '''
        重置噪声，实现重新的进行探索
        todo 用在什么流程上
        '''
        for name, module in self.named_children():
            if 'fc' in name:
                module.reset_noise()

    def both(self, x):
        cat_out, curl_h = self(x)
        probs = self.apply_softmax(cat_out)
        weights = probs * self.supports
        res = weights.sum(dim=2)
        return cat_out, res, curl_h

    def qvals(self, x):
        return self.both(x)[1]

    def apply_softmax(self, t):
        return self.softmax(t.view(-1, self.atoms)).view(t.size())